package essential;

import java.util.ArrayList;
import javax.vecmath.Vector3d;

import primitives.AxisAlignedBoundingBox;
import primitives.Plane;
import primitives.Renderable;

public class KDTree implements KDTreeNode {
	AxisAlignedBoundingBox aabb;
	KDTreeNode child[];

	public KDTreeNode copy() {
		KDTree k = new KDTree();
		k.aabb.setA(aabb.getA());
		k.aabb.setB(aabb.getB());
		k.child[0] = child[0].copy();
		k.child[1] = child[1].copy();
		return k;
	}
	
	public AxisAlignedBoundingBox getAABB() {
		return aabb;
	}
	
	private KDTree() {
		aabb = new AxisAlignedBoundingBox();
		child = new KDTreeNode[2];
	}
	
	private KDTree(KDTreeNode c1, KDTreeNode c2, AxisAlignedBoundingBox aabb) {
		child = new KDTreeNode[2];
		child[0] = c1;
		child[1] = c2;
		this.aabb = aabb;
	}

	enum Side {Positive, Negative, Intersection};
	
	static public KDTreeNode BuildKDTree(ArrayList<Renderable> renderable, AxisAlignedBoundingBox aabb, int depth) {
		Axis axis = chooseAxis(aabb);

		double splitDist = getOptimalDivison(renderable, aabb, axis);
		
		ArrayList<Renderable> posPrimitives = new ArrayList<Renderable>();
		ArrayList<Renderable> negPrimitives = new ArrayList<Renderable>();

		if (Double.isNaN(splitDist)) {
			return new KDTreeLeave(renderable, aabb);
		}

		for(int i= 0; i < renderable.size(); i++) {
			Side side = splitSide(splitDist, renderable.get(i).getAABB(), axis);
			
			if (renderable.get(i) instanceof Plane) {
			} else if (side == Side.Negative) {
				negPrimitives.add(renderable.get(i));
			} else if (side == Side.Positive) {
				posPrimitives.add(renderable.get(i));
			} else {
				negPrimitives.add(renderable.get(i));
				posPrimitives.add(renderable.get(i));
			}
		}

		AxisAlignedBoundingBox negAABB = new AxisAlignedBoundingBox(); 
		AxisAlignedBoundingBox posAABB = new AxisAlignedBoundingBox();

		aabb.split(axis, splitDist, negAABB, posAABB);

//		System.out.println("Spliting in axis " + axis + " at " + splitDist + " depth=" + depth);
//		System.out.println("-Primitives: " + negPrimitives.size());
//		System.out.println("+Primitives: " + posPrimitives.size());
//		System.out.println("-AABB: " + negAABB);
//		System.out.println("+AABB: " + posAABB);

		if (depth > 1) {
			return new KDTree(
					BuildKDTree(negPrimitives, negAABB, depth - 1),
					BuildKDTree(posPrimitives, posAABB, depth - 1),
					aabb);
		} else {
			return new KDTree(
					new KDTreeLeave(negPrimitives, negAABB),
					new KDTreeLeave(posPrimitives, posAABB),
					aabb);
		}
	}

	static Axis chooseAxis(AxisAlignedBoundingBox aabb) {
		Vector3d sizes = new Vector3d(
				aabb.getB().x - aabb.getA().x,
				aabb.getB().y - aabb.getA().y,
				aabb.getB().z - aabb.getA().z);

		return (sizes.x >= sizes.y && sizes.x >= sizes.z)? Axis.X :
			((sizes.y >= sizes.x && sizes.y >= sizes.z)? Axis.Y : Axis.Z);
	}

	static double getOptimalDivison(ArrayList<Renderable> primitives, AxisAlignedBoundingBox aabb, Axis axis) {
		Vector3d sizes = new Vector3d(
				aabb.getB().x - aabb.getA().x,
				aabb.getB().y - aabb.getA().y,
				aabb.getB().z - aabb.getA().z);

		double splits[] = new double[2 * primitives.size()];
		int rcount[] = new int[2 * primitives.size()];
		int lcount[] = new int[2 * primitives.size()];

		//Fill splits with splitting candidates
		for (int i=0; i<primitives.size(); i++) {
			
			
			splits[i*2] = Axis.getComponent(primitives.get(i).getAABB().getA(), axis);
			splits[i*2+1] = Axis.getComponent(primitives.get(i).getAABB().getB(), axis);
		}
		
		// calculate surface area for current node
		double sav = 0.5f / (sizes.x * sizes.y + sizes.x * sizes.z + sizes.y * sizes.z);

		//Calculate cost for not splitting
		double costLeaf = primitives.size() * 1.0f;

		// determine optimal split plane position
		double lowestCost = 10000;
		double bestPos = 0;

		//Get rcount and lcount for each split, and save it if it has the lowest cost
		//Can this(O(N^2) this can be done with lower complexity order)
		//by having the primitives ordered in some way?
		Vector3d rsizes = new Vector3d(sizes);
		Vector3d lsizes = new Vector3d(sizes);
		for (int i=0; i < splits.length; i++) {
			for (int j=0; j < primitives.size(); j++) {
				Side s = splitSide(splits[i], primitives.get(j).getAABB(), axis);
				if (s == Side.Positive || s == Side.Intersection)
					rcount[i]++;
				if (s == Side.Negative || s == Side.Intersection)
					lcount[i]++;
			}

			rsizes = new Vector3d(sizes);
			lsizes = new Vector3d(sizes);
			
			//Calculate the extents of the child boxes
			Axis.setComponent(rsizes, axis, Axis.getComponent(rsizes, axis) - (splits[i] - Axis.getComponent(aabb.getA(), axis)));
			Axis.setComponent(lsizes, axis, splits[i] - Axis.getComponent(aabb.getA(), axis));

			// calculate child node cost
			double rsa = 2 * (rsizes.x * rsizes.y + rsizes.x * rsizes.z + rsizes.y * rsizes.z);
			double lsa = 2 * (lsizes.x * lsizes.y + lsizes.x * lsizes.z + lsizes.y * lsizes.z);
			double splitcost = 0.3f + 1.0f * (rsa * sav * rcount[i] + lsa * sav * lcount[i]);

			// update best cost tracking variables
			if (splitcost < lowestCost) {
				lowestCost = splitcost;
				bestPos = splits[i];
			}
		}
		
		if (lowestCost > costLeaf) {
			return Double.NaN;
		}

		return bestPos;
	}

	static private Side splitSide(double split, AxisAlignedBoundingBox aabb, Axis axis) {
		if ((split == Axis.getComponent(aabb.getA(), axis)) && (split == Axis.getComponent(aabb.getB(), axis)))
			return Side.Intersection;
		else if (split >= Axis.getComponent(aabb.getB(), axis))
			return Side.Negative;
		else if (split <= Axis.getComponent(aabb.getA(), axis))
			return Side.Positive;
		return Side.Intersection;
	}

	public Renderable intersect(Ray r) {
		RayCastResult result1 = child[0].getAABB().intersection(r);
		RayCastResult result2 = child[1].getAABB().intersection(r);
		
		if (result1 == RayCastResult.HIT && result2 == RayCastResult.HIT) {
			Renderable p1 = child[0].intersect(r);
			Renderable p2 = child[1].intersect(r);
			return p2 != null? p2 : p1;

//			[ HERE LAY THE REMAINS OF THE HARDEST BUG EVER		]
//			[ CREATED BY TEH ROW, MAY IT'S 19hs LIFESPAN SERVE	]
//			[ AS AN EXAMPLE FOR FUTURE GENERATIONS. REST IN		]
//			[ PEACE, HONORED AND UNTOUCHED IN THIS HOLY	GRAVE	]
//			[ 													]
//			[ 		"Recursion is a leap of faith" ~ lzabalet	]
//			[													]
//			[					7/5/2010 05:00 ~ 8/5/2010 00:30	]
//			//Copy ray
//			Ray rcopy = new Ray(r);
//			Primitive p1 = child[0].intersect(r);
//			Primitive p2 = child[1].intersect(rcopy);
//
//			if (p1 != null && p2 != null) {
//				if (r.getImpactDist() < rcopy.getImpactDist())
//					return p1;
//				return p2;
//			} else if (p1 != null) {
//				return p1;
//			} else {
//				return p2;
//			}

		} else if (result1 == RayCastResult.HIT) {
			return child[0].intersect(r);
		} else if (result2 == RayCastResult.HIT) {
			return child[1].intersect(r);
		} else {
			return null;
		}
	}
}

class KDTreeLeave implements KDTreeNode {
	ArrayList<Renderable> renderables;
	AxisAlignedBoundingBox aabb;

	public KDTreeNode copy() {
		KDTreeLeave k = new KDTreeLeave();
		k.aabb.setA(aabb.getA());
		k.aabb.setB(aabb.getB());
		
		for(int i=0; i<renderables.size(); i++)
			k.renderables.add(renderables.get(i).copy());
		
		return k;
	}

	public KDTreeLeave() {
		renderables = new ArrayList<Renderable>();
		aabb = new AxisAlignedBoundingBox();
	}
	
	KDTreeLeave(ArrayList<Renderable> renderables, AxisAlignedBoundingBox aabb) {
		this.renderables = renderables;
		this.aabb = aabb;
	}

	public Renderable intersect(Ray r) {
		Renderable result = null;

		// find the nearest intersection
		for (int s = 0; s < renderables.size(); s ++) {
			Renderable rend = renderables.get(s);

			if (r.getLastRenderable() == null || r.getLastRenderable() != rend) {
				if (rend.intersect(r) != RayCastResult.MISS) {
					result = rend;
				}
			}
		}
		return result;
	}
	
	public AxisAlignedBoundingBox getAABB() {
		return aabb;
	}
}
